'''SEGMENTOR_CFG for MAE'''

DATASET_CFG = {
    'type': '',
    'rootdir': '',
    'train': {
        'set': 'train',
        'aug_opts': [
            ('Resize', {'output_size': (2048, 512), 'keep_ratio': True, 'scale_range': (0.5, 2.0)}),
            ('RandomCrop', {'crop_size': (512, 512), 'one_category_max_ratio': 0.75}),
            ('RandomFlip', {'flip_prob': 0.5}),
            ('PhotoMetricDistortion', {}),
            ('Normalize', {'mean': [123.675, 116.28, 103.53], 'std': [58.395, 57.12, 57.375]}),
            ('ToTensor', {}),
            ('Padding', {'output_size': (512, 512), 'data_type': 'tensor'}),
        ],
    },
    'test': {
        'set': 'val',
        'aug_opts': [
            ('Resize', {'output_size': (2048, 512), 'keep_ratio': True, 'scale_range': None}),
            ('Normalize', {'mean': [123.675, 116.28, 103.53], 'std': [58.395, 57.12, 57.375]}),
            ('ToTensor', {}),
        ],
    }
}

DATALOADER_CFG = {
    'train': {
        'batch_size': 16, 'num_workers': 16, 'shuffle': True, 'pin_memory': True, 'drop_last': True,
    },
    'test': {
        'batch_size': 1, 'num_workers': 16, 'shuffle': False, 'pin_memory': True, 'drop_last': False,
    }
}

SCHEDULER_CFG = {
        'type': 'PolyScheduler', 'max_epochs': 10, 'power': 1.0, 'min_lr': 0.0, 
        'warmup_cfg': {'type': 'linear', 'ratio': 1e-6, 'iters': 1500},
}

OPTIMIZER_CFG = {
            'type': 'sgd', 'lr': 1e-4, 'momentum': 0.9, 'weight_decay': 0.0, 
            'params_rules': {'type': 'LayerDecayParamsConstructor', 'num_layers': 12, 'decay_rate': 0.65, 'decay_type': 'layer_wise_vit'},
        }

SEGMENTOR_CFG = {
    'type': 'upernet',
    'num_classes': -1,
    'benchmark': True,
    'align_corners': False,
    'backend': 'nccl',
    # 'norm_cfg': {'type': 'SyncBatchNorm'},
    'norm_cfg': {'type': 'batchnorm2d'},
    # 'act_cfg': {'type': 'ReLU', 'inplace': True},
    'act_cfg': {'type': 'relu'},
    'backbone': {
        'type': 'mae_pretrain_vit_base', 'series': 'mae', 'pretrained': False, 'pretrained_model_path': '',
        'img_size': (512, 512), 'patch_size': 16, 'embed_dims': 768, 'num_layers': 12,
        'num_heads': 12, 'mlp_ratio': 4, 'init_values': 1.0, 'drop_path_rate': 0.1,
        'selected_indices': (0, 1, 2, 3), 'norm_cfg': {'type': 'layernorm', 'epsilon': 1e-6},
    },
    'head': {
        'feature2pyramid': {'embed_dim': 768, 'rescales': [4, 2, 1, 0.5]}, 'in_channels_list': [768, 768, 768, 768], 
        'feats_channels': 768, 'pool_scales': [1, 2, 3, 6], 'dropout': 0.1,
    },
    'auxiliary': {
        'in_channels': 768, 'out_channels': 512, 'dropout': 0.1,
    },
}

LOSSES_CFG = {
        'loss_aux': {'celoss': {'scale_factor': 0.4, 'ignore_index': 255, 'reduction': 'mean'}},
        'loss_cls': {'celoss': {'scale_factor': 1.0, 'ignore_index': 255, 'reduction': 'mean'}},
    }

INFERENCE_CFG = {
        'mode': 'slide',
        'opts': {'cropsize': (512, 512), 'stride': (341, 341)}, 
        'tricks': {
            'multiscale': [1], 'flip': False, 'use_probs_before_resize': False,
        }
    },
# config for common
COMMON_CFG = {
    'work_dir': 'ckpts',
    'logfilepath': '',
    'log_interval_iterations': 50,
    'eval_interval_epochs': 10,
    'save_interval_epochs': 1,
    'resultsavepath': '',
}